﻿#include "modelapplication.h"
#include "ui_modelapplication.h"
#include <QFileDialog>
#include <modelinfodlg.h>
#include <QMessageBox>
#include <QDirIterator>
#include <QDebug>
#include <iostream>
#include "yolo.h"
#include <QSettings>


void ModelApplication::ScanModels()
{
    QSettings *reg = new QSettings("HKEY_CURRENT_USER\\Software\\ImageProcessing", QSettings::NativeFormat);
    m_modelsdir = reg->value("ModelsPath").toString();
    if(!m_modelsdir.trimmed().isEmpty()){
        ui->savepath->setText(m_modelsdir);
    }
    //QMessageBox::information(NULL,QStringLiteral("模型加载"),m_modelsdir);
    //qDebug() << pluginsdir;
    QDirIterator it(m_modelsdir,QDir::Files|QDir::Dirs|QDir::NoDotAndDotDot);	//遍历所有目录和文件
    while (it.hasNext())//存在
    {
        QString name = it.next();//读取
        QFileInfo info(name);
        if (info.isDir())//判断是目录
        {
            //将目录加入运行环境变量
            continue;
        }
        else
        {
            if (info.suffix() == "onnx"){
                QString modelname = info.baseName().toUtf8().data();
                QFile file(info.filePath().replace(".onnx",".txt"));
                if(file.open(QIODevice::ReadOnly)){
                    m_modelinfo[modelname] = file.readAll();
                }else{
                    qDebug() << modelname <<":" <<  QStringLiteral("标签文件加载失败");
                    continue;
                }
                ui->listModels->addItem(modelname);
            }
        }
    }
}

ModelApplication::ModelApplication(QWidget *parent) :
    QWidget(parent),
    ui(new Ui::ModelApplication)
{
    ui->setupUi(this);
    this->setWindowTitle(QStringLiteral("模型工具"));
    ScanModels();
    m_type = MEDIA_IMG;
    m_isPath = false;
}

ModelApplication::~ModelApplication()
{
    if(cap.isOpened())
        cap.release();
    delete ui;
}

void ModelApplication::on_btnadd_clicked()
{
    if(m_modelsdir.trimmed().isEmpty()){
        QMessageBox::information(NULL,QStringLiteral("模型导入"),QStringLiteral("请先选择模型加载路径"));
        return;
    }
    QString url = QFileDialog::getOpenFileName(NULL,QStringLiteral("选择模型"),"./","*.onnx");
    if(!url.isEmpty()){
        //弹出模型信息编辑框
        ModelInfoDlg dlg;
        if(dlg.exec() == QDialog::Accepted){
            //复制保存模型 模型名称 标签信息
            if(m_modelinfo.contains(dlg.m_modelName)){
                QMessageBox::information(NULL,QStringLiteral("模型导入"),QStringLiteral("模型名称已存在，请重新命名"));
                return;
            }
            QString modelpath = m_modelsdir+dlg.m_modelName+".onnx";
            QString labelpath = m_modelsdir+dlg.m_modelName+".txt";
            qDebug() << modelpath;
            if (QFile::exists(modelpath)){
                QFile::remove(modelpath);
            }
            if (QFile::exists(labelpath)){
                QFile::remove(labelpath);
            }
            if(QFile::copy(url,modelpath)){
                qDebug() << QStringLiteral("模型导入成功");
                ui->listModels->addItem(dlg.m_modelName);
                m_modelinfo[dlg.m_modelName] = dlg.m_labelList;
                QFile file(labelpath);
                if(file.open(QIODevice::WriteOnly)){
                    file.write(dlg.m_labelList.toUtf8());
                }
                file.close();
            }else{
                QMessageBox::information(NULL,QStringLiteral("模型导入"),QStringLiteral("模型导入失败"));
            }
        }
    }
}

void ModelApplication::on_btnminus_clicked()
{
    QListWidgetItem *item = ui->listModels->currentItem();
    if(item){
        QString strmodel = item->text();
        if(QMessageBox::question(NULL,QStringLiteral("模型删除"),QStringLiteral("是否要删除'%1'此模型?").arg(strmodel))==QMessageBox::Yes){
            QString modelpath = m_modelsdir+strmodel+".onnx";
            QString labelpath = m_modelsdir+strmodel+".txt";
            qDebug() << modelpath;
            if (QFile::exists(modelpath)){
                QFile::remove(modelpath);
            }
            if (QFile::exists(labelpath)){
                QFile::remove(labelpath);
            }
            m_modelinfo.remove(strmodel);
            ui->listModels->removeItemWidget(item);
            delete item;
        }
    }
}

void ModelApplication::on_listModels_itemDoubleClicked(QListWidgetItem *item)
{
    QMessageBox::information(NULL,QStringLiteral("模型标签"),m_modelinfo[item->text()]);
}

void ModelApplication::on_btnselectimg_clicked()
{
    m_targetPath = QFileDialog::getOpenFileName(NULL,QStringLiteral("打开图片"),"./","Image Files(*.jpg *.png);;");
    if(!m_targetPath.isEmpty()){
        ui->labelpath->setText(m_targetPath);
        m_isPath = false;
        m_type = MEDIA_IMG;
    }
}

void ModelApplication::on_btnselectdir_clicked()
{
    m_targetPath = QFileDialog::getExistingDirectory(this,QStringLiteral("选择目录"),"./",QFileDialog::ShowDirsOnly);
    if(!m_targetPath.isEmpty()){
        ui->labelpath->setText(m_targetPath);
        m_isPath = false;
        m_type = MEDIA_IMG;
    }
}

void ModelApplication::on_listMethods_itemDoubleClicked(QListWidgetItem *item)
{
    if(item->text() == QStringLiteral("图片图像识别")){
        if(m_targetPath.isEmpty()){
            QMessageBox::information(NULL,QStringLiteral("模型识别"),QStringLiteral("请选择图片或文件夹路径"));
            return;
        }
        if(ui->listModels->selectedItems().size()==0){
            QMessageBox::information(NULL,QStringLiteral("识别结果"),QStringLiteral("请选择模型"));
            return;
        }
        QString modename = ui->listModels->currentItem()->text();
        if(!m_modelinfo.contains(modename)){
            QMessageBox::information(NULL,QStringLiteral("识别结果"),QStringLiteral("模型不存在:")+modename);
            return;
        }
        //模型识别处理
        try {
            Yolo yolo;
            dnn::Net net;
            string model_path = (m_modelsdir+modename+".onnx").toUtf8().toStdString();
            if(!yolo.readModel(net,model_path,false)){
                QMessageBox::information(NULL,QStringLiteral("模型加载"),QStringLiteral("模型读取失败:")+modename);
                return;
            }
            if(m_type == MEDIA_IMG){
                cv::Mat image;
                if(m_isPath){
                    QDirIterator it(m_targetPath,QDir::Files|QDir::Dirs|QDir::NoDotAndDotDot); //遍历所有目录和文件
                    while (it.hasNext())//存在
                    {
                        QString name = it.next();//读取
                        QFileInfo info(name);
                        if (info.isDir())//判断是目录
                        {
                            continue;
                        }
                        else
                        {
                            if (info.suffix() == "jpg" || info.suffix() == "bmp" || info.suffix() == "png")
                            {
                                image = cv::imread((m_targetPath+"/"+name).toStdString(), cv::IMREAD_COLOR);  // 读取图片
                            }
                        }
                    }
                }else{
                    //识别图片
                    image = cv::imread(m_targetPath.toStdString(), cv::IMREAD_COLOR);  // 读取图片
                    if (image.empty()){
                        QMessageBox::information(NULL,QStringLiteral("模型识别"),QStringLiteral("图片加载失败"));
                        return;
                    }
                    std::vector<std::string> className = { "A", "B", "C", "D", "E", "F", "G"};
                    // Generate random colors //
                    vector<Scalar> color;
                    srand(time(0));
                    for (int i = 0; i < className.size(); i++) {
                        int b = rand() % 8*32 - 1;
                        int g = rand() % 8*32 - 1;
                        int r = rand() % 8*32 - 1;
                        color.push_back(Scalar(b, g, r));
                    }
                    //color.push_back(Scalar(0, 0, 255));

                    vector<Output> result;
                    if (yolo.Detect(image, net, result,className))// return true Object detected //
                    {
                        yolo.drawPred(image, result, color,className);
                    }
                    else {
                        QMessageBox::information(NULL,QStringLiteral("模型识别"),QStringLiteral("识别失败"));
                        return;
                    }
                    imshow("识别结果",image);
                }
            }
        } catch(const cv::Exception& ex) {
            qDebug() << ex.what() << m_modelsdir+modename+".onnx";
            QMessageBox::information(NULL,QStringLiteral("模型加载"),ex.what());
            return;
        }
    }else if(item->text() == QStringLiteral("摄像头图像识别")){
        if(ui->listModels->selectedItems().size()==0){
            QMessageBox::information(NULL,QStringLiteral("识别结果"),QStringLiteral("请选择模型"));
            return;
        }
        QString modename = ui->listModels->currentItem()->text();
        if(!m_modelinfo.contains(modename)){
            QMessageBox::information(NULL,QStringLiteral("识别结果"),QStringLiteral("模型不存在:")+modename);
            return;
        }
        //模型识别处理
        try {
            Yolo yolo;
            dnn::Net net;
            string model_path = (m_modelsdir+modename+".onnx").toUtf8().toStdString();
            if(!yolo.readModel(net,model_path,true)){
                QMessageBox::information(NULL,QStringLiteral("模型加载"),QStringLiteral("模型读取失败:")+modename);
                return;
            }
            cap = cv::VideoCapture(0);
            cap.set(3,640);
            cap.set(4,480);
            if(!cap.isOpened()){
                QMessageBox::information(NULL,QStringLiteral("模型识别"),QStringLiteral("摄像头打开失败"));
                return;
            }
            cv::Mat frame;
            std::vector<std::string> className = { "A", "B", "C", "D", "E", "F", "G"};
            vector<Scalar> color;
            srand(time(0));
            for (int i = 0; i < className.size(); i++) {
                int b = rand() % 8*32 - 1;
                int g = rand() % 8*32 - 1;
                int r = rand() % 8*32 - 1;
                color.push_back(Scalar(b, g, r));
            }
            cv::namedWindow("camera", WINDOW_NORMAL);
            vector<Output> result;
            while(1){
                cap >> frame;
                if (frame.empty()){
                    QMessageBox::information(NULL,QStringLiteral("模型识别"),QStringLiteral("图片加载失败"));
                    return;
                }
                // Generate random colors //
                //color.push_back(Scalar(0, 0, 255));
                result.clear();
                yolo.Detect(frame, net, result,className);
                if(result.size()){
                    yolo.drawPred(frame, result, color,className);
                }
                /*else {
                    cap.release();
                    QMessageBox::information(NULL,QStringLiteral("模型识别"),QStringLiteral("识别失败"));
                    return;
                }*/
                imshow("camera",frame);
                if (27 == waitKey(30)) {
                    cap.release();
                    destroyWindow("camera");
                    break;
                }
            }
        } catch(const cv::Exception& ex) {
            qDebug() << ex.what() << m_modelsdir+modename+".onnx";
            QMessageBox::information(NULL,QStringLiteral("模型加载"),ex.what());
            return;
        }
    }
}

void ModelApplication::on_btn_path_clicked()
{
    QString dir = QFileDialog::getExistingDirectory(this,QStringLiteral("选择目录"),"./",QFileDialog::ShowDirsOnly);
    if(dir.trimmed().isEmpty()){
        return;
    }
    dir += "/";
    QSettings *reg = new QSettings("HKEY_CURRENT_USER\\Software\\ImageProcessing", QSettings::NativeFormat);
    reg->setValue("ModelsPath",dir);
    m_modelsdir = dir;
    ui->savepath->setText(m_modelsdir);
    ui->listModels->clear();
    ScanModels();
}
